import heapq


class NodesToSearch:
	"""Set of A* nodes to search.
	"""

	def __init__(self, start):
		self.heap = [0]
		self.index = { 0 : [start] }
		self.set = set([start])

	def __contains__(self, item):
		return item in self.set

	def __len__(self):
		return len(self.set)

	def add(self, item, cost):
		"""Add a new node having the associated estimated cost.

		>>> to_search = NodesToSearch('start')
		>>> len(to_search)
		1
		>>> to_search.add('forty', 40)
		>>> len(to_search)
		2
		>>> try:
		... 	to_search.add('start')
		... 	'FAIL'
		... except:
		... 	pass
		"""
		assert item not in self.set
		heapq.heappush(self.heap, cost)
		self.index.setdefault(cost, []).append(item)
		self.set.add(item)

	def pop(self):
		"""Pop and return the node having the lowest estimated cost.

		>>> to_search = NodesToSearch('start')
		>>> to_search.add('one', 1)
		>>> to_search.add('three', 3)
		>>> to_search.add('two', 2)
		>>> to_search.pop()
		'start'
		>>> to_search.pop()
		'one'
		>>> to_search.pop()
		'two'
		>>> to_search.pop()
		'three'
		>>> try:
		... 	to_search.pop()
		... 	'FAIL'
		... except:
		... 	pass
		"""
		item = self.index[heapq.heappop(self.heap)].pop()
		self.set.remove(item)
		return item


class Search:
	"""Class for performing A* searches.
	"""

	def __init__(self, start, goals):
		self.to_search = NodesToSearch(start)
		self.searched = set()
		self.came_from = {}
		self.cost_from_start = { start : 0 }
		self.cost_to_goal = {} # optimistic estimate
		self.goals = goals

	def search(self):
		"""Return the least cost path from start to any of the goals.

		>>> class N(Node):
		...     def __init__(self): self.n = []
		...     def neighbors(self): return self.n
		...     def cost(self, other): return 1
		...     def cost_estimate(self, other): return 1
		>>> n1 = N()
		>>> n2 = N()
		>>> n3 = N()
		>>> n4 = N()
		>>> n1.n = [n2]
		>>> n2.n = [n1, n3]
		>>> n3.n = [n2, n4]
		>>> n4.n = [n3]

		>>> s = Search(n1, [n3])
		>>> path = s.search()
		>>> len(path)
		2
		>>> n2 == path[0]
		True
		>>> n3 == path[1]
		True

		>>> s = Search(n2, [n1, n4])
		>>> path = s.search()
		>>> len(path)
		1
		>>> n1 == path[0]
		True
		"""
		while self.to_search:
			current = self.to_search.pop()

			if current in self.goals:
				return self.path(current)

			self.searched.add(current)

			for neighbor in current.neighbors():
				if neighbor in self.searched: continue

				cost = self.cost_from_start[current] + current.cost(neighbor)
				is_better = False

				if neighbor not in self.to_search:
					self.cost_to_goal[neighbor] = min([neighbor.cost_estimate(g) for g in self.goals])
					self.to_search.add(neighbor, cost + self.cost_to_goal[neighbor])
					is_better = True
				elif cost < self.cost_from_start[neighbor]:
					is_better = True

				if is_better:
					self.came_from[neighbor] = current
					self.cost_from_start[neighbor] = cost

		return None

	def path(self, node):
		"""Return the node sequence path.

		>>> s = Search(None, [])
		>>> s.path('foo')
		[]
		>>> s.came_from['baz'] = 'bar'
		>>> s.came_from['bar'] = 'garply'
		>>> s.path('baz')
		['bar', 'baz']
		>>> s.path('foo')
		[]
		"""
		if node in self.came_from:
			p = self.path(self.came_from[node])
			p.append(node)
			return p
		else:
			return []


class Node:

	def neighbors(self):
		return NotImplementedError

	def cost(self, other):
		return NotImplementedError

	def cost_estimate(self, other):
		return NotImplementedError


if __name__ == '__main__':
	import test
	test.run_doc_tests()

